{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import random\n",
    "from collections import deque\n",
    "import numpy as np\n",
    "import torch.optim as optim\n",
    "import gym"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cartpole_model(observation_space, action_space):\n",
    "    return nn.Sequential(\n",
    "        nn.Linear(observation_space, 24),\n",
    "        nn.ReLU(),\n",
    "        nn.Linear(24, 24),\n",
    "        nn.ReLU(),\n",
    "        nn.Linear(24, action_space)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DQN:\n",
    "    def __init__(self, observation_space, action_space):\n",
    "        self.exploration_rate = MAX_EXPLORE\n",
    "        self.action_space = action_space\n",
    "        self.observation_space = observation_space\n",
    "        self.memory = deque(maxlen=MEMORY_LEN)\n",
    "        \n",
    "        self.target_net = cartpole_model(self.observation_space, self.action_space)\n",
    "        self.policy_net = cartpole_model(self.observation_space, self.action_space)\n",
    "        \n",
    "        self.target_net.load_state_dict(self.policy_net.state_dict())\n",
    "        self.target_net.eval()\n",
    "        \n",
    "        self.criterion = nn.MSELoss()\n",
    "        self.optimizer = optim.Adam(self.policy_net.parameters())\n",
    "        \n",
    "        self.explore_limit = False\n",
    "        \n",
    "    def load_memory(self, state, action, reward, next_state, terminal):\n",
    "        self.memory.append((state, action, reward, next_state, terminal))\n",
    "        \n",
    "    def predict_action(self, state):\n",
    "        random_number = np.random.rand()\n",
    "        \n",
    "        if random_number < self.exploration_rate:\n",
    "            return random.randrange(self.action_space)\n",
    "        \n",
    "        q_values = self.target_net(state).detach().numpy()\n",
    "        return np.argmax(q_values[0])\n",
    "    \n",
    "    def experience_replay(self):\n",
    "        if len(self.memory) < BATCH_SIZE:\n",
    "            return\n",
    "        \n",
    "        batch = random.sample(self.memory, BATCH_SIZE)\n",
    "        \n",
    "        for state, action, reward, next_state, terminal in batch:\n",
    "            q_update = reward\n",
    "            \n",
    "            if not terminal:\n",
    "                    q_update = reward + GAMMA * self.target_net(next_state).max(axis=1)[0]\n",
    "                    \n",
    "            q_values = self.target_net(state)\n",
    "            q_values[0][action] = q_update\n",
    "            \n",
    "            loss = self.criterion(self.policy_net(state), q_values)\n",
    "            self.optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            self.optimizer.step()\n",
    "            \n",
    "        if not self.explore_limit:\n",
    "            self.exploration_rate *= EXPLORE_DECAY\n",
    "            if self.exploration_rate < MIN_EXPLORE:\n",
    "                self.exploration_rate = MIN_EXPLORE\n",
    "                self.explore_limit = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "ENV_NAME = \"CartPole-v1\"\n",
    "BATCH_SIZE = 20\n",
    "GAMMA = 0.95\n",
    "LEARNING_RATE = 0.001\n",
    "MAX_EXPLORE = 1.0\n",
    "MIN_EXPLORE = 0.01\n",
    "EXPLORE_DECAY = 0.995\n",
    "MEMORY_LEN = 1_000_000\n",
    "UPDATE_FREQ = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make(ENV_NAME)\n",
    "observation_space = env.observation_space.shape[0]\n",
    "action_space = env.action_space.n\n",
    "dqn = DQN(observation_space, action_space)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| Run | Exploration Rate | Score |\n",
      "| 001 |       0.7862     |  068  |\n",
      "| 002 |       0.7329     |  015  |\n",
      "| 003 |       0.6369     |  029  |\n",
      "| 004 |       0.5186     |  042  |\n",
      "| 005 |       0.4575     |  026  |\n",
      "| 006 |       0.3157     |  075  |\n",
      "| 007 |       0.2584     |  041  |\n",
      "| 008 |       0.2212     |  032  |\n",
      "| 009 |       0.1670     |  057  |\n",
      "| 010 |       0.1459     |  028  |\n",
      "| 011 |       0.1113     |  055  |\n",
      "| 012 |       0.0788     |  070  |\n",
      "| 013 |       0.0684     |  029  |\n",
      "| 014 |       0.0580     |  034  |\n",
      "| 015 |       0.0390     |  080  |\n",
      "| 016 |       0.0341     |  028  |\n",
      "| 017 |       0.0312     |  019  |\n",
      "| 018 |       0.0210     |  080  |\n",
      "| 019 |       0.0136     |  088  |\n",
      "| 020 |       0.0100     |  071  |\n",
      "| 021 |       0.0100     |  043  |\n",
      "| 022 |       0.0100     |  080  |\n",
      "| 023 |       0.0100     |  103  |\n",
      "| 024 |       0.0100     |  091  |\n",
      "| 025 |       0.0100     |  031  |\n",
      "| 026 |       0.0100     |  099  |\n",
      "| 027 |       0.0100     |  184  |\n",
      "| 028 |       0.0100     |  217  |\n",
      "| 029 |       0.0100     |  029  |\n",
      "| 030 |       0.0100     |  111  |\n",
      "| 031 |       0.0100     |  095  |\n",
      "| 032 |       0.0100     |  135  |\n",
      "| 033 |       0.0100     |  188  |\n",
      "| 034 |       0.0100     |  193  |\n",
      "| 035 |       0.0100     |  188  |\n",
      "| 036 |       0.0100     |  356  |\n",
      "| 037 |       0.0100     |  286  |\n",
      "| 038 |       0.0100     |  225  |\n",
      "| 039 |       0.0100     |  280  |\n",
      "| 040 |       0.0100     |  346  |\n",
      "| 041 |       0.0100     |  316  |\n",
      "| 042 |       0.0100     |  341  |\n",
      "| 043 |       0.0100     |  278  |\n",
      "| 044 |       0.0100     |  500  |\n",
      "| 045 |       0.0100     |  258  |\n",
      "| 046 |       0.0100     |  193  |\n",
      "| 047 |       0.0100     |  347  |\n",
      "| 048 |       0.0100     |  010  |\n",
      "| 049 |       0.0100     |  009  |\n",
      "| 050 |       0.0100     |  009  |\n",
      "| 051 |       0.0100     |  009  |\n",
      "| 052 |       0.0100     |  010  |\n",
      "| 053 |       0.0100     |  009  |\n",
      "| 054 |       0.0100     |  009  |\n",
      "| 055 |       0.0100     |  008  |\n",
      "| 056 |       0.0100     |  009  |\n",
      "| 057 |       0.0100     |  010  |\n",
      "| 058 |       0.0100     |  011  |\n",
      "| 059 |       0.0100     |  183  |\n",
      "| 060 |       0.0100     |  431  |\n",
      "| 061 |       0.0100     |  194  |\n",
      "| 062 |       0.0100     |  219  |\n",
      "| 063 |       0.0100     |  208  |\n",
      "| 064 |       0.0100     |  335  |\n",
      "| 065 |       0.0100     |  170  |\n",
      "| 066 |       0.0100     |  186  |\n",
      "| 067 |       0.0100     |  184  |\n",
      "| 068 |       0.0100     |  168  |\n",
      "| 069 |       0.0100     |  208  |\n",
      "| 070 |       0.0100     |  253  |\n",
      "| 071 |       0.0100     |  192  |\n",
      "| 072 |       0.0100     |  221  |\n",
      "| 073 |       0.0100     |  192  |\n",
      "| 074 |       0.0100     |  206  |\n",
      "| 075 |       0.0100     |  191  |\n",
      "| 076 |       0.0100     |  188  |\n",
      "| 077 |       0.0100     |  246  |\n",
      "| 078 |       0.0100     |  323  |\n",
      "| 079 |       0.0100     |  171  |\n",
      "| 080 |       0.0100     |  231  |\n",
      "| 081 |       0.0100     |  230  |\n",
      "| 082 |       0.0100     |  062  |\n",
      "| 083 |       0.0100     |  194  |\n",
      "| 084 |       0.0100     |  251  |\n",
      "| 085 |       0.0100     |  178  |\n",
      "| 086 |       0.0100     |  171  |\n",
      "| 087 |       0.0100     |  192  |\n",
      "| 088 |       0.0100     |  189  |\n",
      "| 089 |       0.0100     |  183  |\n",
      "| 090 |       0.0100     |  182  |\n",
      "| 091 |       0.0100     |  221  |\n",
      "| 092 |       0.0100     |  178  |\n",
      "| 093 |       0.0100     |  197  |\n",
      "| 094 |       0.0100     |  201  |\n",
      "| 095 |       0.0100     |  268  |\n",
      "| 096 |       0.0100     |  316  |\n",
      "| 097 |       0.0100     |  273  |\n",
      "| 098 |       0.0100     |  173  |\n",
      "| 099 |       0.0100     |  147  |\n",
      "| 100 |       0.0100     |  226  |\n"
     ]
    }
   ],
   "source": [
    "print(f'| Run | Exploration Rate | Score |')\n",
    "steps = 0\n",
    "for i in range(100):\n",
    "    state = env.reset()\n",
    "    state = np.reshape(state, [1, observation_space])\n",
    "    state = torch.from_numpy(state).float()\n",
    "    \n",
    "    score = 0\n",
    "    while True:\n",
    "        steps += 1\n",
    "        score += 1\n",
    "        action = dqn.predict_action(state)\n",
    "        next_state, reward, terminal, info = env.step(action)\n",
    "        \n",
    "        next_state = torch.from_numpy(np.reshape(next_state, [1, observation_space])).float()\n",
    "        dqn.load_memory(state, action, reward, next_state, terminal)\n",
    "        state = next_state\n",
    "        \n",
    "        if terminal:\n",
    "            print(f'| {i+1:03} |       {dqn.exploration_rate:.4f}     |  {score:03}  |')\n",
    "            break\n",
    "        \n",
    "        dqn.experience_replay()\n",
    "        if steps%UPDATE_FREQ == 0:\n",
    "            dqn.target_net.load_state_dict(dqn.policy_net.state_dict())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def play_agent(dqn, env):\n",
    "    observation = env.reset()\n",
    "    total_reward=0\n",
    "    for _ in range(500):\n",
    "        env.render()\n",
    "        observation = torch.tensor(observation).type('torch.FloatTensor').view(1,-1)\n",
    "        q_values = dqn.target_net(observation).detach().numpy()\n",
    "        action = np.argmax(q_values[0])\n",
    "        new_observation, reward, done, _ = env.step(action)\n",
    "        total_reward += reward\n",
    "        observation = new_observation\n",
    "\n",
    "        if(done):\n",
    "            break\n",
    "\n",
    "    env.close()\n",
    "    print(\"Rewards: \",total_reward)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rewards:  160.0\n"
     ]
    }
   ],
   "source": [
    "play_agent(dqn, env)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
